import os
import sys
import matplotlib.pyplot as plt
import numpy as np
import scipy.io.wavfile
from python_ai.common.xcommon import *
from PyCmpltrtok.common import rand_name_on_now
import seaborn as sns

VER = 'v2.0'
BASE_DIR, FILE_NAME = os.path.split(__file__)
SAVE_DIR = os.path.join(BASE_DIR, '_save', FILE_NAME, VER)

sep('Load and emphasize')
path = r'D:\_const\large_data\audio\zsn-stop1.wav'
voice_path = os.path.join(BASE_DIR, path)

sr, y = scipy.io.wavfile.read(voice_path)
print_numpy_ndarray_info(y, 'y')
if len(y.shape) >= 2:
    y = y[:, 0]
print('sr', sr)  # 16000
print('y', y.shape)  # 31744

spr = 2
spc = 1
spn = 0
plt.figure(figsize=[8, 6])

spn += 1
plt.subplot(spr, spc, spn)
plt.title('Sound')
plt.plot(y)
plt.grid()

emphasized = np.append(y[0:1], y[1:] - 0.97 * y[:-1])
print_numpy_ndarray_info(emphasized, 'emphasized')
emphasized_ = emphasized.copy()
xmin = emphasized.min()
xmax = emphasized.max()
emphasized -= xmin
emphasized /= xmax - xmin
emphasized *= 2.
emphasized -= 1.
print_numpy_ndarray_info(emphasized, 'emphasized')
emphasized *= np.iinfo(np.int16).max
emphasized = np.asarray(emphasized, dtype=np.int16)
print_numpy_ndarray_info(emphasized, 'emphasized')

path = os.path.join(SAVE_DIR, 'emphasized' + rand_name_on_now() + '.wav')
print(f'Writing to {path}')
os.makedirs(os.path.split(path)[0], exist_ok=True)
scipy.io.wavfile.write(path, sr, emphasized)
print('Written.')
emphasized = emphasized_

spn += 1
plt.subplot(spr, spc, spn)
plt.title('Emphasized')
plt.plot(emphasized)
plt.grid()

print('Please check and close the plotting window to continue ...')
plt.show()

sep('Frames')
signal_len = len(emphasized)
frame_len = int(np.round(sr / 40))  # 1 / 40 = 1 / 1000 * 25 (25ms)
frame_stride = int(np.round(sr / 100))  # 1 / 100 = 1 / 1000 * 10 (10ms)
n_frames = int(np.ceil(float(np.abs(signal_len - frame_len)) / frame_stride))
print('signal_len', signal_len)  # 31744
print('frame_len', frame_len)  # 400
print('frame_stride', frame_stride)  # 160
print('n_frames', n_frames)  # 196

padded_signal_len = n_frames * frame_stride + frame_len
padding = np.zeros((padded_signal_len - signal_len,), dtype=emphasized.dtype)
padded_signal = np.append(emphasized, padding)
print_numpy_ndarray_info(padded_signal, 'padded_signal')  # (31760, )

idx_col = np.tile(np.arange(0, frame_len), (n_frames, 1))
print(idx_col)
print_numpy_ndarray_info(idx_col, 'idx_col')
idx_row = np.tile(np.arange(0, n_frames * frame_stride, frame_stride).reshape(-1, 1), (1, frame_len))
print(idx_row)
print_numpy_ndarray_info(idx_row, 'idx_row')
idx = np.int64(idx_col + idx_row)
print(idx)
print_numpy_ndarray_info(idx, 'idx')
frames = emphasized[idx]
print_numpy_ndarray_info(frames, 'frames')

spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])

spn += 1
plt.subplot(spr, spc, spn)
plt.title('Frames')
sns.heatmap(frames)

sep('Add windows')
window = np.hamming(frame_len)
print_numpy_ndarray_info(window, 'window')
frames = np.asarray(frames, np.float64)
frames *= window
# frames = np.asarray(frames, np.int16)
print_numpy_ndarray_info(frames, 'frames after window')

spn += 1
plt.subplot(spr, spc, spn)
plt.title('Frames added window')
sns.heatmap(frames)

print('Please check and close the plotting window to continue ...')
plt.show()

plt.plot(window)
plt.title("Hamming window")
plt.grid()
print('Please check and close the plotting window to continue ...')
plt.show()

sep('Fourier transform')
NFFT = 512
rfft = np.fft.rfft(frames, NFFT)
print(rfft)
print_numpy_ndarray_info(rfft, 'rfft')
rfft = np.absolute(rfft)
print(rfft)
print_numpy_ndarray_info(rfft, 'rfft abs')
_, n_spec = rfft.shape  # 257 = 512/2+1

spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])

spn += 1
plt.subplot(spr, spc, spn)
plt.title('rfft')
sns.heatmap(rfft)

sep('Power spectrum')
pow_frames = (rfft ** 2) / NFFT
print_numpy_ndarray_info(pow_frames, 'pow_frames')

spn += 1
plt.subplot(spr, spc, spn)
plt.title('Power spectrum')
sns.heatmap(pow_frames)

print('Please check and close the plotting window to continue ...')
plt.show()

sep('Filter banks')


def hz2mel(hz):
    return 2595 * np.log10(1 + hz / 700)


def mel2hz(mel):
    return (10 ** (mel / 2595) - 1) * 700


n_filters = 40
low_mel_hz = 0
high_hz = sr / 2
high_mel_hz = hz2mel(high_hz)
mel_hz_arr = np.linspace(low_mel_hz, high_mel_hz, n_filters + 2)
hz_arr = mel2hz(mel_hz_arr)
y1y2 = np.tile([[100], [-100]], len(hz_arr))
spr = 2
spc = 2
spn = 0
plt.figure(figsize=[10, 10])
spn += 1
plt.subplot(spr, spc, spn)
plt.plot([hz_arr, hz_arr], y1y2)

bins1 = hz_arr / sr * (NFFT + 1)
print('bins1', bins1)
bins2 = hz_arr / high_hz * n_spec
print('bins2', bins2)
spn += 1
plt.subplot(spr, spc, spn)
plt.plot([bins1, bins1], y1y2)
spn += 1
plt.subplot(spr, spc, spn)
plt.plot([bins2, bins2], y1y2)

bins = bins1

fbanks = np.zeros((n_filters, n_spec), dtype=np.float64)  # (40, 257)
for m in range(1, n_filters + 1):
    f_m_mius = int(bins[m - 1])
    f_m = int(bins[m])
    f_m_plus = int(bins[m + 1])
    for k in range(f_m_mius, f_m):
        fbanks[m - 1, k] = (k - bins[m - 1]) / (bins[m] - bins[m - 1])
    for k in range(f_m, f_m_plus):
        fbanks[m - 1, k] = (bins[m + 1] - k) / (bins[m + 1] - bins[m])
spn += 1
plt.subplot(spr, spc, spn)
for fbank in fbanks:
    plt.plot(fbank)

print('Please check and close the plotting window to continue ...')
plt.show()

sep('Filtered power spectrum')
filtered = np.dot(pow_frames, fbanks.T)  # (196, 257) dot (257, 40) = (196, 40)
print_numpy_ndarray_info(filtered, 'filtered')
filtered = np.where(np.isclose(filtered, 0.), np.finfo(float).eps, filtered)
print_numpy_ndarray_info(filtered, 'filtered')
filtered = 20 * np.log10(filtered)

plt.figure(figsize=[10, 10])
plt.title('Filtered power spectrum')
sns.heatmap(filtered)
plt.show()

